import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def score_correlation(df):
    """
    :param df:
    :return: 返回Pearson相关系数的热力图和系数矩阵
    """

    corr_mtx = df.select_dtypes(include='number').corr()

    # 计算每个科目与其他科目的相关性之和
    sums = {}
    for col in corr_mtx.columns:
        subject = col.split('_')[0]
        if subject not in sums:
            sums[subject] = []
        corrs = corr_mtx[col].values
        corrs = np.delete(corrs, np.where(corr_mtx.columns == col))  # 删除自身相关系数
        sums[subject].append(np.sum(corrs))

    # 计算每个科目的相关性之和均值
    means = {subject: np.mean(sums[subject]) for subject in sums}

    # 排序并输出结果
    sorted_subjects = sorted(means, key=means.get, reverse=True)
    print("各科目重要性排序:")
    for idx, subject in enumerate(sorted_subjects):
        print("{:2}. {} 相关性之和均值: {:.3f}".format(idx + 1, subject, means[subject]))

    # 输出相关系数矩阵
    print('\n相关系数矩阵:')
    print(corr_mtx)

    # 可视化相关矩阵
    plt.rcParams['font.family'] = ['STSong']
    plt.subplots(figsize=(10, 8))
    plt.imshow(corr_mtx, cmap='coolwarm', interpolation='nearest')
    plt.colorbar()
    tick_marks = [i for i in range(len(corr_mtx.columns))]
    plt.xticks(tick_marks, corr_mtx.columns, rotation=45)
    plt.yticks(tick_marks, corr_mtx.columns)
    plt.title('Pearson Correlation Matrix')
    plt.show()

    # 将相关矩阵保存到 Excel 文件
    # corr_mtx.to_excel('pearson_correlation_test001.xlsx', index=None)


def main():
    df = pd.read_excel("../TestExample/test_score001.xlsx")
    score_correlation(df)


if __name__ == "__main__":
    main()
